import numpy as np
from sklearn.cluster import KMeans
from matplotlib import pyplot as plt
from sklearn.metrics import adjusted_rand_score
import random
import copy
import pandas as pd

#使用特征变换
data = np.genfromtxt('dataset_circles.csv',delimiter=',',skip_header=0)
N = len(data)
N_train = int(N*0.8)
N_test = N - N_train

m=[]
np.random.shuffle(data)

for i in range(N):
    m.append([data[i,0]**2+data[i,1]**2,0])
m = np.array(m)

x_train = m[:N_train, :]
y_train = data[:N_train,2]
x_test  = m[N_train:, :]
y_test  = data[N_train:,2]

kmeans = KMeans(n_clusters=2, random_state=0).fit(x_train)

plt.scatter(data[:N_train,0],data[:N_train,1],c=kmeans.labels_)
#plt.scatter(kmeans.cluster_centers_[:,0],kmeans.cluster_centers_[:,1],c='r')
plt.show()
ari_train = adjusted_rand_score(kmeans.predict(x_test),y_test)
print("ari_train = %f" % ari_train)

#不使用特征变换
data = np.genfromtxt('dataset_circles.csv',delimiter=',',skip_header=0)
N = len(data)
N_train = int(N*0.8)
N_test = N - N_train

np.random.shuffle(data)

x_train = data[:N_train, :2]
y_train = data[:N_train,2]
x_test  = data[N_train:, :2]
y_test  = data[N_train:,2]

kmeans = KMeans(n_clusters=2, random_state=0).fit(x_train)

plt.scatter(data[:N_train,0],data[:N_train,1],c=kmeans.labels_)
#plt.scatter(kmeans.cluster_centers_[:,0],kmeans.cluster_centers_[:,1],c='r')
plt.show()
ari_train = adjusted_rand_score(kmeans.predict(x_test),y_test)
print("ari_train = %f" % ari_train)
